#!/usr/bin/env python3
"""
OpenSCAP MCP Server

面向 OpenSCAP 安全合规性扫描工具的命令行操作封装。

MCP tools 列表：
1. oscap_assess_safety_compliance - 评估特定基准的安全性合规性
   参数：xccdf_path (XCCDF 文件路径), profile (配置文件名称), remote_host (远程主机地址), 
        remote_user (远程用户名), remote_password (远程密码)
   返回：统一的命令行执行结果 JSON

2. oscap_view_configuration_files - 查看安全配置文件和基准信息
   参数：content_dir (SCAP 内容目录路径), profile (配置文件名称)
   返回：统一的命令行执行结果 JSON

3. oscap_generate_ansible_fix - 生成 Ansible 修复脚本
   参数：xccdf_path (XCCDF 文件路径), profile (配置文件名称), results_file (结果文件路径)
   返回：统一的命令行执行结果 JSON

4. oscap_scan_remote_vulnerabilities - 扫描远程系统的安全漏洞
   参数：remote_host (远程主机地址), remote_user (远程用户名), remote_password (远程密码), 
        oval_file (OVAL 文件路径)
   返回：统一的命令行执行结果 JSON

5. oscap_scan_local_vulnerabilities - 扫描本地系统的安全漏洞
   参数：oval_file (OVAL 文件路径), report_file (报告文件路径)
   返回：统一的命令行执行结果 JSON

6. oscap_check_remediate_service - 检查 oscap-remediate 服务状态
   参数：无
   返回：统一的命令行执行结果 JSON
"""

import subprocess
import json
from typing import Optional
from mcp.server.fastmcp import FastMCP, Context

mcp = FastMCP("OpenSCAP MCP Server")

@mcp.tool()
def oscap_assess_safety_compliance(
    xccdf_path: str,
    profile: str = "ospp",
    remote_host: Optional[str] = None,
    remote_user: Optional[str] = None,
    remote_password: Optional[str] = None
) -> dict:
    """
    评估特定基准的安全性合规性，支持本地和远程评估
    
    参数:
        xccdf_path: XCCDF 基准文件路径
        profile: 要评估的配置文件名称，默认为 ospp
        remote_host: 远程主机地址（IP或主机名），如未提供则执行本地评估
        remote_user: 远程用户名
        remote_password: 远程密码
    
    返回:
        JSON 结构: {
            "success": bool,      # 命令是否成功执行
            "command": str,       # 执行的完整命令
            "exit_code": int,     # 命令退出码
            "stdout": str,        # 标准输出
            "stderr": str         # 标准错误输出
        }
    """
    try:
        # 本地评估部分
        commands = []
        results = []
        
        # 1. 检查 XCCDF 文档信息
        cmd1 = f"oscap info {xccdf_path}"
        result1 = subprocess.run(cmd1, shell=True, capture_output=True, text=True)
        commands.append(cmd1)
        results.append(result1)
        
        # 2. 检查特定配置文件信息
        cmd2 = f"oscap info --profile {profile} {xccdf_path}"
        result2 = subprocess.run(cmd2, shell=True, capture_output=True, text=True)
        commands.append(cmd2)
        results.append(result2)
        
        # 3. 执行本地评估并生成报告
        cmd3 = f"oscap xccdf eval --report /tmp/report.html --profile {profile} {xccdf_path}"
        result3 = subprocess.run(cmd3, shell=True, capture_output=True, text=True)
        commands.append(cmd3)
        results.append(result3)
        
        # 4. 如果提供了远程主机信息，执行远程评估
        if remote_host and remote_user and remote_password:
            cmd4 = f"expect -c 'set timeout 300; spawn oscap-ssh {remote_user}@{remote_host} 22 xccdf eval --report /tmp/remote_report.html {xccdf_path}; expect {{\\\"*yes/no*\\\" {{send \\\"yes\\\\r\\\"; exp_continue}} \\\"s password: \\\" {{send \\\"{remote_password}\\\\r\\\"; exp_continue}} timeout}}'"
            result4 = subprocess.run(cmd4, shell=True, capture_output=True, text=True)
            commands.append(cmd4)
            results.append(result4)
        
        # 合并所有结果
        combined_stdout = "\n".join([r.stdout for r in results if r.stdout])
        combined_stderr = "\n".join([r.stderr for r in results if r.stderr])
        success = all(r.returncode == 0 for r in results)
        exit_code = 0 if success else 1
        
        return {
            "success": success,
            "command": "; ".join(commands),
            "exit_code": exit_code,
            "stdout": combined_stdout,
            "stderr": combined_stderr
        }
        
    except Exception as e:
        return {
            "success": False,
            "command": "oscap_assess_safety_compliance",
            "exit_code": 1,
            "stdout": "",
            "stderr": f"Error: {str(e)}"
        }

@mcp.tool()
def oscap_view_configuration_files(
    content_dir: str = "/usr/share/xml/scap/ssg/content",
    profile: Optional[str] = None
) -> dict:
    """
    查看安全配置文件和基准信息
    
    参数:
        content_dir: SCAP 内容目录路径，默认为 /usr/share/xml/scap/ssg/content
        profile: 要查询的特定配置文件名称（如 pci-dss）
    
    返回:
        统一的命令行执行结果 JSON
    """
    try:
        commands = []
        results = []
        
        # 1. 列出所有配置文件
        cmd1 = f"for file in {content_dir}/*.xml; do echo \"$file\"; done"
        result1 = subprocess.run(cmd1, shell=True, capture_output=True, text=True)
        commands.append(cmd1)
        results.append(result1)
        
        # 2. 查看默认文档信息
        default_file = f"{content_dir}/ssg-ol7-ds.xml"
        cmd2 = f"oscap info {default_file}"
        result2 = subprocess.run(cmd2, shell=True, capture_output=True, text=True)
        commands.append(cmd2)
        results.append(result2)
        
        # 3. 如果指定了配置文件，查看特定配置信息
        if profile:
            # 查找最新的 SL 文件
            cmd3 = f"find {content_dir} -name 'ssg-sl[0-9]*-xccdf.xml' -print0 | sort -z -V | tail -z -n 1 | tr -d '\\0'"
            result3 = subprocess.run(cmd3, shell=True, capture_output=True, text=True)
            latest_file = result3.stdout.strip()
            
            if latest_file:
                cmd4 = f"oscap info --profile {profile} \"{latest_file}\""
                result4 = subprocess.run(cmd4, shell=True, capture_output=True, text=True)
                commands.append(cmd4)
                results.append(result4)
        
        # 合并结果
        combined_stdout = "\n".join([r.stdout for r in results if r.stdout])
        combined_stderr = "\n".join([r.stderr for r in results if r.stderr])
        success = all(r.returncode == 0 for r in results)
        exit_code = 0 if success else 1
        
        return {
            "success": success,
            "command": "; ".join(commands),
            "exit_code": exit_code,
            "stdout": combined_stdout,
            "stderr": combined_stderr
        }
        
    except Exception as e:
        return {
            "success": False,
            "command": "oscap_view_configuration_files",
            "exit_code": 1,
            "stdout": "",
            "stderr": f"Error: {str(e)}"
        }

@mcp.tool()
def oscap_generate_ansible_fix(
    xccdf_path: str,
    profile: str = "ospp",
    results_file: str = "ospp-results.xml"
) -> dict:
    """
    生成 Ansible 修复脚本使系统符合特定标准
    
    参数:
        xccdf_path: XCCDF 基准文件路径
        profile: 要评估的配置文件名称，默认为 ospp
        results_file: 评估结果文件路径
    
    返回:
        统一的命令行执行结果 JSON
    """
    try:
        commands = []
        results = []
        
        # 1. 执行评估并保存结果
        cmd1 = f"oscap xccdf eval --profile {profile} --results {results_file} {xccdf_path}"
        result1 = subprocess.run(cmd1, shell=True, capture_output=True, text=True)
        commands.append(cmd1)
        results.append(result1)
        
        # 2. 生成 Ansible 修复配置
        fix_file = "ospp-remediations.yml"
        cmd2 = f"oscap xccdf generate fix --fix-type ansible --output {fix_file} {results_file}"
        result2 = subprocess.run(cmd2, shell=True, capture_output=True, text=True)
        commands.append(cmd2)
        results.append(result2)
        
        # 3. 验证生成的修复文件
        cmd3 = f"grep host {fix_file}"
        result3 = subprocess.run(cmd3, shell=True, capture_output=True, text=True)
        commands.append(cmd3)
        results.append(result3)
        
        # 合并结果
        combined_stdout = "\n".join([r.stdout for r in results if r.stdout])
        combined_stderr = "\n".join([r.stderr for r in results if r.stderr])
        success = all(r.returncode == 0 for r in results)
        exit_code = 0 if success else 1
        
        return {
            "success": success,
            "command": "; ".join(commands),
            "exit_code": exit_code,
            "stdout": combined_stdout,
            "stderr": combined_stderr
        }
        
    except Exception as e:
        return {
            "success": False,
            "command": "oscap_generate_ansible_fix",
            "exit_code": 1,
            "stdout": "",
            "stderr": f"Error: {str(e)}"
        }

@mcp.tool()
def oscap_scan_remote_vulnerabilities(
    remote_host: str,
    remote_user: str,
    remote_password: str,
    oval_file: str = "/usr/share/xml/scap/ssg/content/ssg-ol7-oval.xml"
) -> dict:
    """
    扫描远程系统的安全漏洞
    
    参数:
        remote_host: 远程主机地址
        remote_user: 远程用户名
        remote_password: 远程密码
        oval_file: OVAL 漏洞定义文件路径
    
    返回:
        统一的命令行执行结果 JSON
    """
    try:
        cmd = f"expect -c 'set timeout 3600; spawn oscap-ssh {remote_user}@{remote_host} 22 oval eval --report /tmp/remote-vulnerability.html {oval_file}; expect {{\\\"*yes/no*\\\" {{send \\\"yes\\\\r\\\"; exp_continue}} \\\"s password: \\\" {{send \\\"{remote_password}\\\\r\\\"; exp_continue}} timeout}}'"
        
        result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
        
        return {
            "success": result.returncode == 0,
            "command": cmd,
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
        
    except Exception as e:
        return {
            "success": False,
            "command": "oscap_scan_remote_vulnerabilities",
            "exit_code": 1,
            "stdout": "",
            "stderr": f"Error: {str(e)}"
        }

@mcp.tool()
def oscap_scan_local_vulnerabilities(
    oval_file: str = "/usr/share/xml/scap/ssg/content/ssg-ol7-oval.xml",
    report_file: str = "vulnerability.html"
) -> dict:
    """
    扫描本地系统的安全漏洞
    
    参数:
        oval_file: OVAL 漏洞定义文件路径
        report_file: 生成的报告文件路径
    
    返回:
        统一的命令行执行结果 JSON
    """
    try:
        cmd = f"oscap oval eval --report {report_file} {oval_file}"
        
        result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
        
        return {
            "success": result.returncode == 0,
            "command": cmd,
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
        
    except Exception as e:
        return {
            "success": False,
            "command": "oscap_scan_local_vulnerabilities",
            "exit_code": 1,
            "stdout": "",
            "stderr": f"Error: {str(e)}"
        }

@mcp.tool()
def oscap_check_remediate_service() -> dict:
    """
    检查 oscap-remediate 服务状态
    
    参数:
        无
    
    返回:
        统一的命令行执行结果 JSON
    """
    try:
        cmd = "systemctl status oscap-remediate.service"
        
        result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
        
        return {
            "success": result.returncode == 0,
            "command": cmd,
            "exit_code": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr
        }
        
    except Exception as e:
        return {
            "success": False,
            "command": "oscap_check_remediate_service",
            "exit_code": 1,
            "stdout": "",
            "stderr": f"Error: {str(e)}"
        }

if __name__ == "__main__":
    mcp.run()